import pandas as pd
import numpy as np
from scipy.optimize import least_squares
import matplotlib.pyplot as plt

# 读取xlsx文件
file_path = '41467_2023_44344_MOESM3_ESM.xlsx'  # 请确保文件路径正确
df = pd.read_excel(file_path, skiprows=1)  # 跳过前两行

# 将数据转换为数值类型
df['Unnamed: 0'] = pd.to_numeric(df['Unnamed: 0'], errors='coerce')
df['PDA'] = pd.to_numeric(df['PDA'], errors='coerce')
df['PDA.1'] = pd.to_numeric(df['PDA.1'], errors='coerce')

# 删除包含NaN值的行
df.dropna(inplace=True)

# 打印前5行数据
print(df.head())

# 打印字段（列名）
print(df.columns)

def realimag(array):
    return np.array([(x.real, -x.imag) for x in array])

def func(x, p):
    s, u, sigma, t = p  # s: static dielectric constant; u: optical dielectric constant; sigma: conductivity; t: relaxation time
    d = complex(0, 1)
    o = 8.854187817 * 10**(-12)
    return realimag(u + (s - u) / (1 + np.dot(d, np.dot(x, t))) - np.dot(d, np.divide(sigma, np.dot(x, o))))

def conduct_loss_result(x, p):
    s, u, sigma, t = p
    o = 8.854187817 * 10**(-12)
    return np.divide(sigma, np.dot(x, o))

def relax_loss_result(x, p):
    s, u, sigma, t = p
    return np.dot(np.divide(s - u, (1 + np.dot(x, t) ** 2)), np.dot(x, t))

def residuals(p, y, x):
    return (realimag(np.array(y)) - func(x, p)).flatten()

p0 = [50, 5, 1, 10**-11]
fcost = 0
fsig = 0  # fitting conductivity
ft = 0  # fitting relaxation time
fplsq = []
fs = 0  # fitting static dielectric constant
fu = 0  # fitting optical dielectric constant
fconduct_loss = 0  # fitting conducting loss
frelax_loss = 0  # fitting polarization loss

# 假设我们只处理第一列数据（例如，PDA）
xdata = df['Unnamed: 0'].values
ydata_1 = df['PDA.1'].values  # experimental imaginary permittivity
ydata_2 = df['PDA'].values  # experimental real permittivity

# 存储拟合结果
fit_results = []

for i in range(0, len(xdata), 10):
    end = i + 10
    start = end - 10
    if end > len(xdata):
        end = len(xdata)
    xdata_segment = xdata[start:end]  # experimental angular frequency
    ydata_1_segment = ydata_1[start:end]  # experimental imaginary permittivity
    ydata_2_segment = ydata_2[start:end]  # experimental real permittivity
    ydata = []
    ydata1_mean = np.mean(ydata_1_segment)
    ydata2_mean = np.mean(ydata_2_segment)
    for j in range(len(ydata_1_segment)):
        ydata.append(complex(ydata_2_segment[j], -ydata_1_segment[j]))
    plsq = least_squares(residuals, p0, bounds=([0, 0, 0, 0], [100, 100, 100, 10**-10]), args=(ydata, xdata_segment), max_nfev=100000)
    fplsq.append(plsq)
    fs = plsq.x[0]
    fu = plsq.x[1]
    fsig = plsq.x[2]
    ft = plsq.x[3]
    fcost = plsq.cost
    fconduct_loss = np.mean(conduct_loss_result(xdata_segment, plsq.x))
    frelax_loss = np.mean(relax_loss_result(xdata_segment, plsq.x))
    o = 8.854187817 * 10**(-12)  # 定义真空介电常数
    fydata2 = fu + np.divide((fs - fu), (1 + (np.dot(xdata_segment, ft) ** 2)))
    fydata1 = np.dot(np.divide((fs - fu), (1 + (np.dot(xdata_segment, ft) ** 2))), np.dot(xdata_segment, ft)) + np.divide(fsig, np.dot(xdata_segment, o))
    fit_results.append((xdata_segment, fydata1, fydata2))
    print(f'{xdata_segment[0]}, {fconduct_loss}, {frelax_loss}, {np.mean(fydata1)}, {ydata1_mean}, {fs}, {fu}, {fsig}, {ft}')

# 绘制太赫兹吸收谱图
plt.figure(figsize=(10, 6))
for xdata_segment, fydata1, fydata2 in fit_results:
    plt.plot(xdata_segment, fydata1, label='Imaginary Part')
    plt.plot(xdata_segment, fydata2, label='Real Part')
plt.xlabel('Frequency (THz)')
plt.ylabel('Permittivity')
plt.title('Terahertz (THz) Spectrum of PDA@OHG')
plt.legend()
plt.grid(True)
plt.show()

# 绘制吸收稳定性图
plt.figure(figsize=(10, 6))
for xdata_segment, fydata1, fydata2 in fit_results:
    plt.plot(xdata_segment, fydata1, label='Imaginary Part')
plt.xlabel('Frequency (THz)')
plt.ylabel('Imaginary Permittivity')
plt.title('Absorbing Stability of PDA@OHG')
plt.legend()
plt.grid(True)
plt.show()

# 计算反射损耗
Z0 = 377  # 自由空间的特征阻抗
reflection_loss_results = []

for xdata_segment, fydata1, fydata2 in fit_results:
    epsilon_r = fydata2 + 1j * fydata1
    Z_in = Z0 * np.sqrt(1 / epsilon_r)
    reflection_loss = 20 * np.log10(np.abs((Z_in - Z0) / (Z_in + Z0)))
    reflection_loss_results.append((xdata_segment, reflection_loss))

# 绘制反射损耗图
plt.figure(figsize=(10, 6))
for xdata_segment, reflection_loss in reflection_loss_results:
    plt.plot(xdata_segment, reflection_loss, label='Reflection Loss')
plt.xlabel('Frequency (THz)')
plt.ylabel('Reflection Loss (dB)')
plt.title('Reflection Loss of PDA@OHG')
plt.legend()
plt.grid(True)
plt.show()
